import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
import os
import random


plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
plt.rc('font', size=10)
plt.rc('font', family='SimHei')


def set_seed(seed=308):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)


set_seed(308)  # 固定随机种子，方便结果复现


# 通用常数
interval = 1  # 投弹时间间隔
v_down = 3  # 云团下降速度
r_effective = 10  # 有效遮蔽半径
t_effective = 20  # 有效遮蔽时间
v_missile = 300  # 导弹速度
r_target = 7  # 目标半径
H_target = 10  # 目标高度
target = np.array([0, 200, 0])  # 目标底部圆心坐标
M1 = np.array([20000, 0, 2000])
M2 = np.array([19000, 600, 2100])
M3 = np.array([18000, -600, 1900])
FY1 = np.array([17800, 0, 1800])
FY2 = np.array([12000, 1400, 1400])
FY3 = np.array([6000, -3000, 700])
FY4 = np.array([11000, 2000, 1800])
FY5 = np.array([13000, -2000, 1300])
v_range = np.array([70, 140])  # 无人机速度范围
g = 9.8  # 重力加速度
# target_point = np.array([
#     [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, 0],
#     [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, 0],
#     [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, 0],
#     [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, 0],
#     [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target],
#     [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target],
#     [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target],
#     [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target],
#     [(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target / 2],
#     [-(1 / 2) ** 0.5 * r_target, (1 / 2) ** 0.5 * r_target + 200, H_target / 2],
#     [(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target / 2],
#     [-(1 / 2) ** 0.5 * r_target, -(1 / 2) ** 0.5 * r_target + 200, H_target / 2],
# ])
target_point = np.array([
    [0, 200, 0]
])
t_total = np.linalg.norm(M1) / 300  # 从0开始到导弹击中目标的总用时
precision = 0.01  # 时间精度


# 传入导弹的初始坐标（也即可确定其飞行方向）和经过的时间，返回其在给定时间下的坐标
def get_M_coordinate(coordinate_0, t, v=None):
    if not v:
        v = -coordinate_0 / np.linalg.norm(coordinate_0) * v_missile
    else:
        v = v / np.linalg.norm(v) * v_missile
    return coordinate_0 + v * t


# 传入第i架无人机的坐标和经过的时间，返回它投射的第j枚烟雾弹中心的坐标
def get_FY_coordinate(coordinate_0, t, i, v_1, t_1, t_2, theta, j):
    t_drop = t_1[i][j]
    t_explosion = t_2[i][j]
    if t < t_drop + t_explosion or t > t_drop + t_explosion + t_effective:
        return [None, None, None]
    else:
        v = np.array([np.cos(theta[i]), np.sin(theta[i]), 0]) * v_1[i]
        delta_t = t - t_drop - t_explosion
        coordinate = coordinate_0[i] + v * (t_drop + t_explosion)
        coordinate[2] -= 1 / 2 * g * (t_explosion ** 2)
        coordinate[2] -= delta_t * v_down
        return coordinate


def collision_detection(coordinate_missile, coordinate_0, t, v_1, t_1, t_2, theta):
    global target_point, r_effective
    count_collision = 0
    for target in target_point:
        for i in range(5):
            flag = False
            for j in range(3):
                coordinate_bomb = get_FY_coordinate(coordinate_0, t, i, v_1, t_1, t_2, theta, j)
                if coordinate_bomb[0] == None:
                    continue
                l_1 = coordinate_bomb - target
                l_2 = coordinate_missile - target
                molecule = np.cross(l_1, l_2)
                denominator = np.linalg.norm(l_2)
                d = np.linalg.norm(molecule) / denominator
                if d < r_effective:
                    missile_to_target = -coordinate_missile  # 由导弹指向目标
                    missile_to_bomb = coordinate_bomb - coordinate_missile  # 由导弹指向烟幕弹
                    cos = np.dot(missile_to_target, missile_to_bomb)
                    if cos > 0 or np.linalg.norm(missile_to_bomb) < r_effective:  # 余弦值大于0，是锐角，或者导弹和烟幕弹之间的距离小于有效半径
                        count_collision += 1
                        flag = True
                        break
            if flag:
                break
    return count_collision


# 飞行方向、飞行速度、烟幕干扰弹投放点、烟幕干扰弹起爆点
def fun(theta, v_1, t_1, t_2):
    t_list = np.arange(0, t_total, precision)
    count = 0
    for t in t_list:
        for M in [M2, M1, M3]:
            if collision_detection(get_M_coordinate(M, t), [FY1, FY2, FY3, FY4, FY5], t, v_1, t_1, t_2, theta) == len(target_point):
                count += 1
    return count


# 在SA函数外部初始化记录容器
history = {
    'best_P': [],
    'current_P': [],
    'k': [[], [], [], [], [], [], [], []],
}


def plot_SA_history():
    plt.figure(figsize=(15, 8))

    # 子图1：损失函数变化
    plt.subplot(2, 3, 1)
    plt.plot(history['best_P'], 'r-', label='Best P')
    plt.plot(history['current_P'], 'b--', alpha=0.5, label='Current P')
    plt.xlabel('Iteration')
    plt.ylabel('P')
    plt.title('P变化曲线')
    plt.legend()
    plt.grid(True, which="both", ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 2)
    plt.plot(history['k'][0], label="FY1飞行方向角")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('飞行方向角变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 3)
    plt.plot(history['k'][1], label="无人机速度大小")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('飞行速度变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 4)
    plt.plot(history['k'][2], label="投弹时间")
    plt.plot(history['k'][3], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹1时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 5)
    plt.plot(history['k'][4], label="投弹时间")
    plt.plot(history['k'][5], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹2时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    # 子图3：k参数演化
    plt.subplot(2, 3, 6)
    plt.plot(history['k'][6], label="投弹时间")
    plt.plot(history['k'][7], label="爆炸时间")
    plt.xlabel('Iteration')
    plt.ylabel('value')
    plt.title('烟幕弹3时间参数变化曲线')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, ls="--")

    plt.tight_layout()
    plt.show()

# 定义变量范围
#   theta    v t1_1 t2_1 t1_2 t2_2 t1_3 t2_3  theta   v t1_1 t2_1 t1_2 t2_2 t1_3 t2_3  theta  v  t1_1 t2_1 t1_2 t2_2 t1_3 t2_3 theta   v t1_1 t2_1 t1_2 t2_2 t1_3 t2_3 theta   v  t1_1 t2_1 t1_2 t2_2 t1_3 t2_3
lb = [3  , 100,   0,   3,   1,   1,   1,  1,     5,  95,  9, 4.5,  1,   1,  1,   1,    1.2,  75,  30,  3, 1,    1,  1,   1,     3, 120,  30,  15,   1,   1,   1,  1,   1,     100,  10,   3,   1,    1,   1,   1]
ub = [3.5, 120,   2,   4,   5,   5,   5,  5,   5.5, 105, 12, 5.5,  5,   5,  5,   5,    1.6,  85,  40,  4, 5,    5,  5,   5,     4, 130,  40,  20,   5,   5,   5,  5,   2,     140,  20,   5,   5,    5,   5,   5]


# 定义目标函数
def objective(x):
    theta_0, v_1_0, t_1_0_0, t_2_0_0, t_1_0_1, t_2_0_1, t_1_0_2, t_2_0_2, theta_1, v_1_1, t_1_1_0, t_2_1_0, t_1_1_1, t_2_1_1, t_1_1_2, t_2_1_2, theta_2, v_1_2, t_1_2_0, t_2_2_0, t_1_2_1, t_2_2_1, t_1_2_2, t_2_2_2, theta_3, v_1_3, t_1_3_0, t_2_3_0, t_1_3_1, t_2_3_1, t_1_3_2, t_2_3_2, theta_4, v_1_4, t_1_4_0, t_2_4_0, t_1_4_1, t_2_4_1, t_1_4_2, t_2_4_2 = x
    theta = [theta_0, theta_1, theta_2, theta_3, theta_4]
    v_1 = [v_1_0, v_1_1, v_1_2, v_1_3, v_1_4]
    t_1 = [[t_1_0_0, t_1_0_0 + t_1_0_1, t_1_0_0 + t_1_0_1 + t_1_0_2], [t_1_1_0, t_1_1_0 + t_1_1_1, t_1_1_0 + t_1_1_1 + t_1_1_2], [t_1_2_0, t_1_2_0 + t_1_2_1, t_1_2_0 + t_1_2_1 + t_1_2_2], [t_1_3_0, t_1_3_0 + t_1_3_1, t_1_3_0 + t_1_3_1 + t_1_3_2], [t_1_4_0, t_1_4_0 + t_1_4_1, t_1_4_0 + t_1_4_1 + t_1_4_2]]
    t_2 = [[t_2_0_0, t_2_0_1, t_2_0_2], [t_2_1_0, t_2_1_1, t_2_1_2], [t_2_2_0, t_2_2_1, t_2_2_2], [t_2_3_0, t_2_3_1, t_2_3_2], [t_2_4_0, t_2_4_1, t_2_4_2]]
    return -fun(theta, v_1, t_1, t_2)  # DE默认最小化，所以取负

# 运行PSO
xopt, fopt = pso(objective, lb, ub, swarmsize=10, maxiter=1)

# 提取最佳解
theta_0, v_1_0, t_1_0_0, t_2_0_0, t_1_0_1, t_2_0_1, t_1_0_2, t_2_0_2, theta_1, v_1_1, t_1_1_0, t_2_1_0, t_1_1_1, t_2_1_1, t_1_1_2, t_2_1_2, theta_2, v_1_2, t_1_2_0, t_2_2_0, t_1_2_1, t_2_2_1, t_1_2_2, t_2_2_2, theta_3, v_1_3, t_1_3_0, t_2_3_0, t_1_3_1, t_2_3_1, t_1_3_2, t_2_3_2, theta_4, v_1_4, t_1_4_0, t_2_4_0, t_1_4_1, t_2_4_1, t_1_4_2, t_2_4_2 = xopt
theta = [theta_0, theta_1, theta_2, theta_3, theta_4]
v_1 = [v_1_0, v_1_1, v_1_2, v_1_3, v_1_4]
t_1 = [[t_1_0_0, t_1_0_0 + t_1_0_1, t_1_0_0 + t_1_0_1 + t_1_0_2],
       [t_1_1_0, t_1_1_0 + t_1_1_1, t_1_1_0 + t_1_1_1 + t_1_1_2],
       [t_1_2_0, t_1_2_0 + t_1_2_1, t_1_2_0 + t_1_2_1 + t_1_2_2],
       [t_1_3_0, t_1_3_0 + t_1_3_1, t_1_3_0 + t_1_3_1 + t_1_3_2],
       [t_1_4_0, t_1_4_0 + t_1_4_1, t_1_4_0 + t_1_4_1 + t_1_4_2]]
t_2 = [[t_2_0_0, t_2_0_1, t_2_0_2], [t_2_1_0, t_2_1_1, t_2_1_2], [t_2_2_0, t_2_2_1, t_2_2_2],
       [t_2_3_0, t_2_3_1, t_2_3_2], [t_2_4_0, t_2_4_1, t_2_4_2]]

Pb = -fopt  # 恢复目标值
print(f"最佳解: P = {Pb}\ntheta = {theta}\nv_1 = {v_1}\nt_1 = {t_1}\nt_2 = {t_2}")
